#define LOG_TAG "udp"
#include "esp_rom_md5.h"
#include "esp_random.h"
#include "esp_ota_ops.h"
#include "esp_mac.h"
#include "esp_netif_ip_addr.h"
#include "esp_log.h"
#ifdef CONFIG_DRIVER_USE_OTA
#include "ota.h"
#endif
#include "udp.h"
#include "wifi.h"
#include "file_tool.h"
static const char *TAG = "UDP/Server";

#ifdef CONFIG_SPIRAM
#define _malloc_(size) heap_caps_malloc(size, MALLOC_CAP_SPIRAM)
#else
#define _malloc_(size) malloc(size)
#endif

#define UDP_MAGIC 0x10090126
// #define ESP_LOGE(TAG, ...)
// #define ESP_LOGI(TAG, ...)
// #define ESP_LOGW(TAG, ...)

#define MULTICAST_LOOPBACK 0
#define MULTICAST_TTL 128

#ifdef CONFIG_SPIRAM_ALLOW_STACK_EXTERNAL_MEMORY
#undef CONFIG_SPIRAM_ALLOW_STACK_EXTERNAL_MEMORY
#endif

typedef struct udp_packet
{
    struct udp_packet *next;
    uint32_t source_ip;
    uint16_t port;
    uint8_t len;
    uint8_t read;
    uint8_t buff[UDP_PACKET_BUFFER_SIZE];
} udp_packet, *udp_packet_t;

typedef struct udp_multicast
{
    uint32_t MAGIC_START;
    TaskHandle_t task;
    char addr[16];
    uint16_t port;
    udp_packet_t list;
    uint32_t last_ip;
    uint16_t last_port;
    int8_t packet_len;
    udp_packet_t send_list;
    int8_t send_len;
    int sock;
    struct sockaddr_in hints;
    esp_app_desc_t run_app_info;
    esp_ip4_addr_t local_ip;
    uint8_t mac[6];
    uint8_t buf[128];
    char token[32];
#if CONFIG_SPIRAM_ALLOW_STACK_EXTERNAL_MEMORY
    uint32_t THREAD_STACK_TOP_MASK;
    StackType_t thread_stack_buffer[3 * 1024];
    uint32_t THREAD_END_MASK;
    StaticTask_t *thread_task;
#endif
    uint32_t MAGIC_END;
} udp_multicast, *udp_multicast_t;

udp_multicast_t _udp_ = 0;

/**
 *      md5校验
 *      mac     : 设备地址
 *      key     : 临时授权密钥
 *      md5_str : 需要比较的md5字符串
 *      成功返回 1
 *      失败返回 0
 */
static uint8_t md5_check(uint8_t *mac, uint32_t key, char *md5_str)
{
    if (!md5_str)
        return 0;
    uint8_t md5[16] = {0};
    char temp[50] = {0};
    char *str = NULL;
    sprintf(temp, MACSTR "|%lu", MAC2STR(mac), key);

    md5_context_t context;
    esp_rom_md5_init(&context);
    esp_rom_md5_update(&context, (unsigned char *)temp, strlen(temp));
    esp_rom_md5_final(md5, &context);
    str = get_hex_str(md5, 16);
    return !strncmp(str, md5_str, 32);
}

bool udp_overflow_check(udp_multicast_t udp)
{
    if (!udp)
    {
        ESP_LOGE(TAG, "!!! Udp Overflow !!!");
        return true;
    }
    if (udp->MAGIC_START != UDP_MAGIC)
    {
        ESP_LOGE(TAG, "!!! Udp stack Overflow !!!");
        return true;
    }
    if (udp->MAGIC_END != UDP_MAGIC)
    {
        ESP_LOGE(TAG, "!!! Udp End Overflow !!!");
        return true;
    }
    return false;
}

/* Add a socket, to the IPV4 multicast group */
int socket_add_ipv4_multicast_group(udp_multicast_t config, int sock, bool assign_source_if)
{
    struct ip_mreq imreq = {0};
    struct in_addr iaddr = {0};
    int err = 0;
    // Configure source interface
    imreq.imr_interface.s_addr = IPADDR_ANY;
    // Configure multicast address to listen to
    err = inet_aton(config->addr, &imreq.imr_multiaddr.s_addr);
    if (err != 1)
    {
        ESP_LOGE(TAG, "Configured IPV4 multicast address '%s' is invalid", config->addr);
        // Errors in the return value have to be negative
        err = -1;
        goto err;
    }
    ESP_LOGI(TAG, "Configured IPV4 Multicast address %s", inet_ntoa(imreq.imr_multiaddr.s_addr));
    if (!IP_MULTICAST(ntohl(imreq.imr_multiaddr.s_addr)))
        ESP_LOGW(TAG, "Configured IPV4 multicast address '%s' is not a valid multicast address. This will probably not work", config->addr);

    if (assign_source_if)
    {
        // Assign the IPv4 multicast source interface, via its IP
        // (only necessary if this socket is IPV4 only)
        err = setsockopt(sock, IPPROTO_IP, IP_MULTICAST_IF, &iaddr,
                         sizeof(struct in_addr));
        if (err < 0)
        {
            ESP_LOGE(TAG, "Failed to set IP_MULTICAST_IF. Error %d", errno);
            goto err;
        }
    }

    err = setsockopt(sock, IPPROTO_IP, IP_ADD_MEMBERSHIP,
                     &imreq, sizeof(struct ip_mreq));
    if (err < 0)
    {
        ESP_LOGE(TAG, "Failed to set IP_ADD_MEMBERSHIP. Error %d", errno);
        goto err;
    }

err:
    return err;
}

int create_multicast_ipv4_socket(udp_multicast_t config)
{
    struct sockaddr_in saddr = {0};
    int sock = -1;
    int err = 0;

    sock = socket(PF_INET, SOCK_DGRAM, IPPROTO_IP);
    if (sock < 0)
    {
        ESP_LOGE(TAG, "Failed to create socket. Error %d", errno);
        return -1;
    }

    // Bind the socket to any address
    saddr.sin_family = PF_INET;
    saddr.sin_port = htons(config->port);
    saddr.sin_addr.s_addr = htonl(INADDR_ANY);
    err = bind(sock, (struct sockaddr *)&saddr, sizeof(struct sockaddr_in));
    if (err < 0)
    {
        ESP_LOGE(TAG, "Failed to bind socket. Error %d", errno);
        goto err;
    }

    // // Assign multicast TTL (set separately from normal interface TTL)
    uint8_t ttl = MULTICAST_TTL;
    setsockopt(sock, IPPROTO_IP, IP_MULTICAST_TTL, &ttl, sizeof(uint8_t));
    if (err < 0)
    {
        ESP_LOGE(TAG, "Failed to set IP_MULTICAST_TTL. Error %d", errno);
        goto err;
    }

#if MULTICAST_LOOPBACK
    // select whether multicast traffic should be received by this device, too
    // (if setsockopt() is not called, the default is no)
    uint8_t loopback_val = MULTICAST_LOOPBACK;
    err = setsockopt(sock, IPPROTO_IP, IP_MULTICAST_LOOP,
                     &loopback_val, sizeof(uint8_t));
    if (err < 0)
    {
        ESP_LOGE(TAG, "Failed to set IP_MULTICAST_LOOP. Error %d", errno);
        goto err;
    }
#endif

    // this is also a listening socket, so add it to the multicast
    // group for listening...
    err = socket_add_ipv4_multicast_group(config, sock, true);
    if (err < 0)
    {
        goto err;
    }

    // All set, socket is configured for sending and receiving
    return sock;

err:
    close(sock);
    return -1;
}

void mcast_example_task(void *pvParameters)
{
    udp_multicast_t config = (udp_multicast_t)pvParameters;
    {
        const esp_partition_t *run_app_partition = esp_ota_get_running_partition();
        esp_ota_get_partition_description(run_app_partition, &config->run_app_info);
    }
    config->local_ip.addr = wifi_get_localIP();
    esp_base_mac_addr_get(config->mac);
    uint32_t key = 0;
    while (1)
    {
        int sock;
        sock = create_multicast_ipv4_socket(config);
        config->sock = sock;
        if (sock < 0)
            ESP_LOGE(TAG, "Failed to create IPv4 multicast socket");

        if (sock < 0)
        {
            // Nothing to do!
            vTaskDelay(pdMS_TO_TICKS(100));
            continue;
        }

        // set destination multicast addresses for sending from these sockets
        struct sockaddr_in sdestv4 = {
            .sin_family = PF_INET,
            .sin_port = htons(config->port),
        };
        ESP_LOGI(TAG, "Listen port : %d", config->port);
        // We know this inet_aton will pass because we did it above already
        inet_aton(config->addr, &sdestv4.sin_addr.s_addr);

        // udp.send((uint8_t *)"TTG start.", 11);
        // ESP_LOGI(TAG, "udp send 'TTG start' message");

        // Loop waiting for UDP received, and sending UDP packets if we don't
        // see any.
        int err = 1;
        while (err > 0 || sock >= 0)
        {
            udp_overflow_check(config);
            struct timeval tv = {
                .tv_sec = 1,
                .tv_usec = 0,
            };
            fd_set rfds;
            FD_ZERO(&rfds);
            FD_SET(sock, &rfds);

            int s = select(sock + 1, &rfds, NULL, NULL, &tv);
            if (s < 0)
            {
                ESP_LOGE(TAG, "Select failed: errno %d", errno);
                err = -1;
                break;
            }
            else if (s > 0)
            {
                if (FD_ISSET(sock, &rfds))
                {
                    // Incoming datagram received
                    memset(config->buf, 0, sizeof(config->buf));
                    struct sockaddr_storage raddr; // Large enough for both IPv4 or IPv6
                    socklen_t socklen = sizeof(raddr);
                    int len = recvfrom(sock, config->buf, UDP_PACKET_BUFFER_SIZE, 0,
                                       (struct sockaddr *)&raddr, &socklen);
                    if (len < 0)
                    {
                        ESP_LOGE(TAG, "multicast recvfrom failed: errno %d", errno);
                        err = -1;
                        break;
                    }

                    int send_len = 0;
                    if (strstr((char *)config->buf, "TTGGetInfo"))
                    {
                        send_len = snprintf((char *)config->buf, sizeof(config->buf),
                                            IPSTR " " MACSTR " %s %s %s",
                                            IP2STR(&config->local_ip),
                                            MAC2STR(config->mac),
                                            config->run_app_info.project_name,
                                            config->run_app_info.version,
                                            config->run_app_info.time);
                        uint32_t ip = ((struct sockaddr_in *)&raddr)->sin_addr.s_addr;
                        uint8_t *ip_p = (uint8_t *)&ip;
                        uint16_t port = ntohs(((struct sockaddr_in *)&raddr)->sin_port);
                        // ESP_LOGI(TAG, IPSTR ":%d",
                        //          ip_p[0], ip_p[1], ip_p[2], ip_p[3],
                        //          port);
                    }
                    else if (strstr((char *)config->buf, "TTGGetKey")) /* TTGGetKey mac:\x64\xb7\x08\x32\xbf\xf4 */
                    {
                        char *macs = strstr((char *)config->buf, "mac:");
                        if (macs)
                        {
                            macs += 4;
                            char *macs_end = (char *)config->buf + len;
                            while (macs < macs_end)
                            {
                                if (memcmp(macs, config->mac, 6) == 0)
                                {
                                    key = esp_random();
                                    send_len = snprintf((char *)config->buf, sizeof(config->buf), MACSTR " %lu", MAC2STR(config->mac), key);
                                }
                                else
                                {
                                    ESP_LOGW(TAG, "NOT FOUND mac: " MACSTR " " MACSTR, MAC2STR(config->mac), MAC2STR(macs));
                                }
                                macs += 6;
                            }
                        }
                        if (!send_len)
                            continue;
                    }
                    else if (strstr((char *)config->buf, "TTGSetToken"))
                    {
                        /* "TTGSetToken md5=95e7fcc2881cb6338e5b951d96e128fd token=8471" */
                        if (key)
                        {
                            char *md5 = strstr((char *)config->buf, "md5=");
                            char *token = strstr((char *)config->buf, "token=");
                            if (md5 && token)
                            {
                                md5 += 4;
                                token += 6;
                                if (md5_check(config->mac, key, md5))
                                {
                                    strncpy(config->token, token, sizeof(config->token));
                                    send_len = snprintf((char *)config->buf, sizeof(config->buf), "OK");
                                }
                            }
                        }
                        if (!send_len)
                            send_len = snprintf((char *)config->buf, sizeof(config->buf), "Reconfigure interrupt");
                    }
#ifdef CONFIG_DRIVER_USE_OTA
                    else if (strstr((char *)config->buf, "TTGProgrees"))
                    {
                        /* TTGProgrees mac:\x64\xb7\x08\x32\xbf\xf4 */
                        char *macs = strstr((char *)config->buf, "mac:");
                        if (macs)
                        {
                            macs += 4;
                            char *macs_end = (char *)config->buf + len;
                            while (macs < macs_end)
                            {
                                if (memcmp(macs, config->mac, 6) == 0)
                                {
                                    if (ota_default_handle->status != OTA_STATUS_RUNNING)
                                        send_len = snprintf((char *)config->buf, sizeof(config->buf), "progress -%d", ota_default_handle->status);
                                    else
                                        send_len = snprintf((char *)config->buf, sizeof(config->buf), "progress %d%%", ota_default_handle->progress);
                                }
                                macs += 6;
                            }
                        }
                        if (!send_len)
                            continue;
                    }
                    else if (strstr((char *)config->buf, "TTGOTA"))
                    {
                        do
                        {
                            char *type_str = (char *)config->buf;
#ifdef ENV_PRODUCTION
                            if (key == 0)
                                goto Permission_denied;
                            /* "TTGOTA token=8471 app 192.168.3.247 80 /ota/bin" */
                            char *token = strstr(type_str, "token=");
                            if (token)
                            {
                                token += 6;
                                type_str = strchr(token, ' ');
                                *type_str = '\0';
                                ESP_LOGI(TAG, "token=%s %s", config->token, token);
                                if (strncmp(config->token, token, sizeof(config->token)))
                                    goto Permission_denied;
                            }
                            else
#endif
                                type_str = strchr(type_str, ' ');
                            /* "TTGOTA app 192.168.3.247 80 /ota/bin" */
                            uint8_t type = 0;
                            if (!type_str)
                                goto Insufficient_parameters;

                            type_str++;
                            if (!strncmp(type_str, "app", 3))
                                type = OTA_TYPE_APP;
                            else if (!strncmp(type_str, "www", 3))
                            {
                                type = OTA_TYPE_WWW;
                                ota_set_par_name(ota_default_handle, "www");
                            }
                            else
                                goto Insufficient_parameters;

                            char *ip = strchr(type_str, ' ');
                            if (!ip)
                                goto Insufficient_parameters;
                            *ip = '\0';
                            ip++;
                            char *port = strchr(ip, ' ');
                            if (!port)
                                goto Insufficient_parameters;
                            *port = '\0';
                            port++;
                            char *path = strchr(port, ' ');
                            if (!path)
                                goto Insufficient_parameters;
                            *path = '\0';
                            path++;
                            if (path[strlen(path) - 1] == '\n')
                                path[strlen(path) - 1] = '\0';
                            ESP_LOGI(TAG, "ip:%s port:%s path:%s type:%s(%d)", ip, port, path, type_str, type);
                            ota_set_server(ota_default_handle, ip, port, path, type);

                            send_len = snprintf((char *)config->buf, sizeof(config->buf), "OK");
                            break;
#ifdef ENV_PRODUCTION
                        Permission_denied:
                            send_len = snprintf((char *)config->buf, sizeof(config->buf), "Permission denied");
                            break;
#endif
                        Insufficient_parameters:
                            send_len = snprintf((char *)config->buf, sizeof(config->buf), "Insufficient parameters");
                            break;
                        } while (0);
                    }
#endif
                    if (send_len)
                    {
                        // elog_hexdump("", 16, config->buf, send_len);
                        sendto(sock, config->buf, send_len, 0, (struct sockaddr *)&raddr, sizeof(struct sockaddr));
                        continue;
                    }

                    udp_packet_t packet = 0;
                    uint8_t get_last = 0;
                    if (UDP_MAX_PACKET > config->packet_len)
                    {
                        packet = _malloc_(sizeof(udp_packet));
                        if (!packet)
                        {
                            if (!config->list)
                            {
                                ESP_LOGE(TAG, "No mem");
                                break;
                            }
                            else
                                get_last = 1;
                        }
                    }
                    else
                        get_last = 1;

                    if (get_last)
                    {
                        packet = config->list;
                        udp_packet_t last = config->list;
                        while (packet->next)
                            last = packet, packet = packet->next;
                        memset(packet->buff, 0, packet->len);
                        last->next = 0;
                    }

                    memcpy(packet->buff, config->buf, len);

                    packet->source_ip = ((struct sockaddr_in *)&raddr)->sin_addr.s_addr;
                    packet->port = ntohs(((struct sockaddr_in *)&raddr)->sin_port);
                    packet->len = len;
                    if (config->list)
                        packet->next = config->list;
                    else
                        packet->next = 0;
                    config->list = packet;
                    if (!get_last)
                        config->packet_len++;

                    config->last_ip = packet->source_ip;
                    config->last_port = packet->port;

                    char raddr_name[32] = {0};
                    if (raddr.ss_family == PF_INET)
                    {
                        inet_ntoa_r(((struct sockaddr_in *)&raddr)->sin_addr,
                                    raddr_name, sizeof(raddr_name) - 1);
                    }
                    ESP_LOGI(TAG, "received %d bytes from %s:%d", len, raddr_name, packet->port);
                }
            }
            else
            {
                // int len = snprintf((char *)config->buf, sizeof(config->buf), IPSTR " " MACSTR " %s %s", IP2STR(&config->local_ip), MAC2STR(config->mac), config->run_app_info.project_name, config->run_app_info.version);
                // udp.send(config->buf, len);
            }
        }

        ESP_LOGW(TAG, "Shutting down socket and restarting...");
        shutdown(sock, 0);
        close(sock);
        config->sock = -1;
        vTaskDelay(1000 / portTICK_PERIOD_MS);
    }
}

udp_multicast_t udp_create_multicast(char *addr, uint16_t port)
{
    if (!_udp_)
    {
        _udp_ = (udp_multicast_t)_malloc_(sizeof(udp_multicast));
        if (!_udp_)
        {
            ESP_LOGE(TAG, "No men to Create udp multicast config");
            return 0;
        }
        memset(_udp_, 0, sizeof(udp_multicast));
        _udp_->MAGIC_START = _udp_->MAGIC_END = UDP_MAGIC;
    }
    memcpy(_udp_->addr, addr, strlen(addr));
    _udp_->port = port;
    if (_udp_->task)
        vTaskDelete(_udp_->task);

#if CONFIG_SPIRAM_ALLOW_STACK_EXTERNAL_MEMORY
    static StaticTask_t _udp_task = {0};
    _udp_->thread_task = &_udp_task;
    _udp_->THREAD_STACK_TOP_MASK = _udp_->THREAD_END_MASK = UDP_MAGIC;
    _udp_->task = xTaskCreateStaticPinnedToCore(mcast_example_task, "mcast_udp", sizeof(_udp_->thread_stack_buffer), _udp_, 5, _udp_->thread_stack_buffer, _udp_->thread_task, 1);
    if (!_udp_->task)
#else
    if (xTaskCreatePinnedToCore(&mcast_example_task, "mcast_udp", 3072, _udp_, 5, &_udp_->task, portNUM_PROCESSORS - 1) != pdTRUE)
#endif
        ESP_LOGE(TAG, "Failed to create mcast_udp thread");
    return _udp_;
}

uint8_t udp_parsePacket(udp_multicast_t config)
{
    if (!config || !config->list)
        return 0;
    udp_packet_t packet = config->list;
    while (packet->next)
        packet = packet->next;
    return packet->len;
}

uint8_t udp_read(udp_multicast_t config, uint8_t *buff, uint8_t len)
{
    if (!config || !config->list)
        return 0;
    udp_packet_t packet = config->list, last = config->list;
    while (packet->next)
        last = packet, packet = packet->next;
    uint8_t size = len;
    if ((packet->len - packet->read) < len)
        size = packet->len - packet->read;
    memcpy(buff, &packet->buff[packet->read], size);
    packet->read += size;
    if (packet->len == packet->read)
    {
        if (packet == config->list)
            config->list = 0;
        else
            last->next = 0;
        // config->last_ip = packet->source_ip;
        // config->last_port = packet->port;
        config->packet_len--;
        free(packet);
    }
    return size;
}
uint8_t udp_create_send(udp_multicast_t config, char *addr, uint16_t port, uint8_t *buff, uint8_t len)
{
    if (!config || !config->sock)
        return 0;
    memset(&config->hints, 0, sizeof(config->hints));
    config->hints.sin_family = AF_INET;
    config->hints.sin_addr.s_addr = inet_addr(addr);
    config->hints.sin_port = htons(port);
    int err = sendto(config->sock, buff, len, 0, (struct sockaddr *)&config->hints, sizeof(config->hints));
    if (err < 0)
        goto send_error;
    return 1;
send_error:
    shutdown(config->sock, 0);
    close(config->sock);
    config->sock = -1;
    return 0;
}

udp_multicast_t udp_config = 0;
void _udp_multicast(char *addr, uint16_t port)
{
    if (!udp_config)
    {
        udp_config = udp_create_multicast(addr, port);
        if (!udp_config)
        {
            ESP_LOGE(TAG, "Failed create multicast");
            return;
        }
    }
}
uint8_t _udp_parsePacket(void)
{
    return udp_parsePacket(udp_config);
}
uint8_t _udp_read(uint8_t *buff, uint8_t len)
{
    return udp_read(udp_config, buff, len);
}
uint8_t _udp_send_multicast(uint8_t *buff, uint8_t len)
{
    if (!udp_config)
        return 0;
    return udp_create_send(udp_config, udp_config->addr, udp_config->port, buff, len);
}
uint8_t _udp_send_to(char *addr, uint16_t port, uint8_t *buff, uint8_t len)
{
    return udp_create_send(udp_config, addr, port, buff, len);
}
uint8_t _udp_send_to_cur(uint8_t *buff, uint8_t len)
{
    if (!udp_config)
        return 0;
    udp_packet_t packet = udp_config->list;
    if (packet)
    {
        while (packet->next)
            packet = packet->next;
        char ip[] = "255.255.255.255";
        esp_ip4_addr_t source_ip = {.addr = packet->source_ip};
        sprintf(ip, IPSTR, IP2STR(&source_ip));
        return udp_create_send(udp_config, ip, packet->port, buff, len);
    }
    else if (udp_config->last_ip)
    {
        char ip[] = "255.255.255.255";
        esp_ip4_addr_t source_ip = {.addr = udp_config->last_ip};
        sprintf(ip, IPSTR, IP2STR(&source_ip));
        return udp_create_send(udp_config, ip, udp_config->last_port, buff, len);
    }
    else
        return udp_create_send(udp_config, udp_config->addr, udp_config->port, buff, len);
}

struct udp udp = {
    .multicast = _udp_multicast,
    .parsePacket = _udp_parsePacket,
    .read = _udp_read,
    .send = _udp_send_multicast,
    .send_to = _udp_send_to,
    .send_to_cur = _udp_send_to_cur};
